#!/scratch/project/qaafi-cnafs/jingping/mamba/bin/python
import sys,os
from subprocess import call
from Config import *
import argparse



def get_head_cmd(threads, Mem):
    cmd_sh = f'''#!/bin/bash --login
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task={threads}
#SBATCH --mem={Mem}G
#SBATCH --job-name=autoslum
#SBATCH --time=40:00:00
#SBATCH --partition=general
#SBATCH --account=a_qaafi_genomics
#SBATCH -o allhic.slurm.out
#SBATCH -e allhic.slurm.error\n
'''
    return cmd_sh

class Scaffold(object):
    def __init__(self, reference, fq1, fq2, threads, outdir, aligntools, agpfile, bwaopt, samtoolsopt):
        self.reference = reference
        self.ref = reference
        self.fq1 = fq1
        self.fq2 = fq2
        self.threads = threads
        self.agpfile = agpfile
        self.outdir = outdir
        self.aligntools = aligntools
        self.bwaopt = bwaopt
        self.samtoolsopt = samtoolsopt
    #bwa_options default is "-SP5M", samtools options default is " -hF 268 "
    def alignment(self, aligntools, bwa_options, samtools_options):
        ref = self.reference
        dir_ref = os.path.dirname(os.path.abspath(ref))
        ref_list = os.listdir(dir_ref)
        outdir = self.outdir
        threads = self.threads
        fq1 = self.fq1
        fq2 = self.fq2
        agpfile = self.agpfile
        bwa_index = 0
        bowtie2_index = 0
        chromap_index = 0
        chromapindex = f"{dir_ref}/chromap.index"
        bowtie2index = "{dir_ref}/contigs"
        cmd = ""
        for i in ref_list:
            if i.endswith("bwt"):
                bwa_index = 1
            if i.endswith("bt2"):
                bowtie2_index = 1
                prefix = i.split(".")[0]
                bowtie2index = f"{dir_ref}/{prefix}"
            if i.endswith("index"):
                chromap_index = 1
                chromapindex = f"{dir_ref}/{i}"
        if aligntools == "bwa":
            if bwa_index == 0:
                cmd += f"{bwa} index {ref}\n"
            cmd += f"{bwa} mem {bwa_options} -t {threads} {ref} {fq1} {fq2} | {samtools} view {samtools_options} - | {samtools} sort  -@  {threads}  -o  {outdir}/sample.bwa_mem.bam  -T  {outdir}/tmp.align\n"

            # cmd += f"{samtools} sort -n -@ 12 -o {outdir}/aligned.sort_name.bam {outdir}/sample.bwa_mem.bam\n"
            # cmd += f"{samtools} index  {outdir}/aligned.sort_name.bam\n"
            # cmd += f"{matlock} bam2 juicer  {outdir}/aligned.sort_name.bam  {outdir}/merged_nodups.txt\n"
            # cmd += f"{}"
            bam = f"{outdir}/sample.bwa_mem.bam"
        elif aligntools == "bowtie2":
            if bowtie2_index == 0:
                cmd += f"{bowtie2_index} bowtie2-build {ref} {bowtie2index} -p {threads}\n"
            cmd += f"{bowtie2} -p {threads} -x {bowtie2index} -1 {fq1} -2 {fq2} > {outdir}/align.bowtie2.sam\n"
            cmd += f"{samtools} view -S -b {outdir}/align.bowtie2.sam > {outdir}/align.bowtie2.bam\n"
            cmd += f"rm {outdir}/align.bowtie2.sam\n"
            bam = f"{outdir}/align.bowtie2.bam"
        elif aligntools == "chromap":
            if chromap_index == 0:
                cmd += f"{chromap} -i -r {ref} -o {chromapindex}"
            cmd += f"{chromap} --preset hic -x {chromapindex} -r {ref} -1 {fq1} -2 {fq2} -o {outdir}/align.chromap.sam --SAM -t {threads}\n"
            cmd += f"{chromap} --preset hic -x {chromapindex} -r {ref} -1 {fq1} -2 {fq2} -o {outdir}/align.chromap.pairs -t {threads}\n"
            cmd += f"{samtools} view -S -b {outdir}/align.chromap.sam > {outdir}/align.chromap.bam\n"
            cmd += f"rm {outdir}/align.chromap.sam\n"
            bam = f"{outdir}/align.chromap.bam"
            pair = f"{outdir}/align.chromap.pairs"
            self.pair = pair
        self.bam = bam
        return cmd

    def bam2juicer(self):
        cmd = self.alignment(self.aligntools, self.bwaopt, self.samtoolsopt)
        outdir = self.outdir
        agp = self.agpfile
        threads = self.threads
        bam = self.bam
        cmd += f"{samtools} sort -n -@ {threads} -o {outdir}/aligned.sort_name.bam {bam}\n"
        cmd += f"{samtools} index {outdir}/aligned.sort_name.bam\n"
        cmd += f"{matlock} bam2 juicer {outdir}/aligned.sort_name.bam merged_nodups.txt\n"
        cmd += f"{py2} {agp2assamble} {agp} {outdir}/groups.assembly\n"
        cmd += f"bash {assamble2hic} -q 1 -p true  {outdir}/groups.assembly  {outdir}/merged_nodups.txt\n"
        return cmd
    
    def pair2juicer(self):
        pair = self.pair
        tmp = '''{if($6!="+") $6=16; else $6=0; if ($7!="+") $7=16;else $7=0}'''
        tmp1 = '''{print $6, $2, $3, 0, $7, $4, $5, 1, "1 - - 1  - - -"}'''
        tmp2 = '''{print $7, $4, $5, 0, $6, $2, $3, 1, "1 - - 1 - - -" }'''
        cmd = f"grep -Ev '#' {pair} | awk '{tmp}$2<$4{tmp1} $4<$2{tmp2}' {pair} > sample.links.txt\n"
        cmd += f"{py2} {agp2assamble} {agpfile} {outdir}/groups.assembly\n"
        cmd += f"bash {assamble2hic} -q 1 -p true  {outdir}/groups.assembly  {outdir}/merged_nodups.txt\n"
        return cmd
    
    def salsa2(self, REs):
        alncmd = self.alignment(self.aligntools, self.bwaopt, self.samtoolsopt)
        cmd = alncmd
        bam = self.bam
        cmd += f"bamToBed -i {bam} > {self.outdir}/align.bed\n"
        cmd += f"sort -k 4 {self.outdir}/align.bed > {self.outdir}/tmp && mv {self.outdir}/tmp {self.outdir}/align.bed\n"
        if not os.path.exists(self.ref + ".fai"):
            cmd += f"samtools faidx {self.ref}\n"
        cmd += "module load salsa/2.3\n"
        cmd += f"run_pipeline.py -a {self.ref} -l {self.ref}.fai -b {self.outdir}/align.bed -o {self.outdir}/salsa_scaffolds -e {REs} --clean yes --exp 877\n"
        return cmd
    
    def haphic(self, RE, haphic_options):
        ref = self.reference
        dir_ref = os.path.dirname(os.path.abspath(ref))
        ref_list = os.listdir(dir_ref)
        outdir = self.outdir
        threads = self.threads
        fq1 = self.fq1
        fq2 = self.fq2
        bwaopt = self.bwaopt
        samtoolsopt = self.samtoolsopt
        cmd += f"{bwa} mem {bwaopt} -t {threads} {ref} {fq1} {fq2} | {samblaster} | {samtools} view -@ {threads} {samtoolsopt} -o  {outdir}/sample.bwa_mem.bam\n"
        cmd += "source activate haphic\n"
        cmd += f"{haphic_bin}/utils/filter_bam.py {self.bam} 1 --NM 3 --threads 96 | samtools view - -b -@ 96 -o {self.outdir}/HiC.filtered.bam\n"
        #cmd += f"{haphic_bin}/haphic pipeline {self.ref} {self.bam} {nchrs} --RE {REs} --threads {threads} --correct_nrounds 2 --max_inflation 5\n"
        cmd += "cd {self.outdir}\n"
        cmd += f"{haphic_bin}/haphic pipeline {self.ref} {self.bam} --RE {RE} {haphic_options} --threads {threads}\n" 
        return cmd
    
        
        
class Assamble(object):
    def __init__(self, fq, outdir, threads, readstype, HiCfq):
        self.ontfq = fq
        self.readstype = readstype
        if readstype == "HiFi":
            self.hififq = fq
        else:
            self.ontfq = fq
        self.outdir = os.path.abspath(outdir)
        if not os.path.exists(self.outdir):
            os.mkdir(self.outdir)
        self.threads = threads
        if HiCfq != None:
            self.HiCfq = HiCfq
            
    def fastqc(self):
        outdir = self.outdir
        trimeddir = f"{outdir}/trimmed_fq/"
        if not os.path.exists(trimeddir):
            os.mkdir(trimeddir)
        fq1, fq2 = [i.strip() for i in self.HiCfq.split(",")]
        self.qcfq1 = f"{trimeddir}/trimmed.R1.fastq.gz"
        self.qcfq2 = f"{trimeddir}/trimmed.R2.fastq.gz"
        cmd = f"{fastp} -q 20 -l 75 -w 8 -i {fq1} -I {fq2} {self.qcfq1} {self.qcfq2}\n"
        return cmd
    
    def hifiasmble(self):
        threads = self.threads
        outdir = f"{self.outdir}/hifiasm_result/"
        if not os.path.exists(outdir):
            os.mkdir(outdir)
        fq1 = self.qcfq1
        fq2 = self.qcfq2
        hififqtmp = self.hififq.split(",")
        if len(hififqtmp) > 1:
            hififq = ",".join(hififqtmp)
        else:
            hififq = self.hififq
        cmd = f"{hifiasm} -t {threads} -l0 --n-hap 3 --h1 {fq1} --h2 {fq2} {hififq}\n"
        return cmd
    
    
    def busco(self):
        pass
        

class AnnotGenome(object):
    def __init__(self, genome, RNAseqFqdir, outdir):
        self.genome = genome
        self.RNAseqFqdir = RNAseqFqdir
        self.outdir = outdir
    
    def RepeatMasker(self):
        cmd = "module load repeatmodeler/2.0.4--pl5321hdfd78af_0/module\n"
        prefix = self.genome.split("/")[-1].split(".")[0]
        cmd += f"BuildDatabase -name {self.outdir}/{prefix} {self.geneome}\n"
        cmd += f"RepeatModeler -database {self.outdir}/{prefix} -threads 20 -LTRStruct >& {self.outdir}/run.out\n"
        
        
        

def assamble(args):
    outdir = args.outdir
    threads = args.threads
    fq = args.fq
    Hicfq = args.Hicfq
    datatype = args.datatype
    asm = Assamble(fq, outdir, threads, datatype, Hicfq)
    qccmd = asm.fastqc()
    asmcmd = asm.hifiasm()
    Mem = args.Mem
    cmd = qccmd + asmcmd
    head = get_head_cmd(threads, Mem)
    cmdout = open(f"{outdir}/assemble.slum", "w")
    cmdout.write(head)
    cmdout.write(cmd)
    cmdout.close()
    runcmd = f"cd {outdir} && sbatch assemble.slum"
    call(runcmd, shell=True)
    
def scaffold(args):
    ref = args.ref
    outdir = args.outdir
    fq1 = args.fq1
    fq2 = args.fq2
    threads = args.threads
    aligntools = args.mt
    mem = args.Mem
    agp = args.agp
    RE = args.RE
    actions = args.action
    bwaopt = args.bwaopt
    samtoolsopt = args.samtoolsopt
    head = get_head_cmd(threads, mem)
    haphicopt = args.haphicopt
    scaffold = Scaffold(ref, fq1, fq2, threads, outdir, aligntools, agp, bwaopt, samtoolsopt)
    if actions == "salsa2":
        cmd = scaffold.salsa2(RE)
    if actions == "bam2juicer":
        cmd = scaffold.bam2juicer(RE)
    if actions == "haphic":
        cmd = scaffold.haphic(RE, haphicopt)
    cmdout = open(f"{outdir}/scaffold.slum", "w")
    cmdout.write(head)
    cmdout.write(cmd)
    cmdout.close()
    cmdrun = "cd {outdir} && sbatch scaffold.slum" 
    
def alignment(args):
    ref = args.ref
    fq1 = args.fq1
    fq2 = args.fq2
    outdir = args.outdir
    threads = args.threads
    mem = args.Mem
    datatype = args.datatype
    head = get_head_cmd(threads, mem)
    tmpopt = '''{sum+=$3;rd+=$4; cov += $5; dp+=$7*$5 }END{print "Total: "sum"\\nCovbase: "cov"\\ncoverage: "cov/sum"\\nmeandepth: "dp/sum}'''
    if datatype == "ngs":
        cmd = f"{bwa} mem -t {threads} {ref} {fq1} {fq2} |samtools view -bh - | {samtools} sort - -@ {threads} -o {outdir}/aln.sort.bam\n"
    else:
        if datatype == "hifi":
            ax = "map-pb"
        elif datatype == "ont":
            ax = "map-ont"
        cmd = f"{minimap2} -t 80 -ax {ax} --secondary=no {ref} {fq1} | samtools view -bh - | samtools sort -@ 80 -o aln.sort.bam - \n"
    cmd += f"{samtools} index {outdir}/aln.sort.bam\n"
    cmd += f"{samtools} flagstat {outdir}/aln.sort.bam > {outdir}/aln.flagstats\n"
    cmd += f"{samtools} coverage {outdir}/aln.sort.bam > {outdir}/aln.coverage\n"
    cmd += f"awk 'NR > 1 {tmpopt}' {outdir}/aln.coverage > aln.coverage.all\n"
    cmdout = open(f"{outdir}/align.slum", "w")
    cmdout.write(head)
    cmdout.write(cmd)
    cmdout.close()
    cmdrun = f"cd {outdir} && sbatch align.slum"
    os.system(cmdrun)

    
    
def switch_error(args):
    ref = args.ref
    ngsfq1 = args.fq1
    ngsfq2 = args.fq2
    longfq = args.longfq
    outdir = args.outdir
    hicfq1 = args.hicfq1
    hicfq2 = args.hicfq2
    threads = args.threads
    mem = args.Mem
    

def structure_annot(args):
    
    
    
    
    


    
    
    
    
    
    
    
    

def main():
    parser = argparse.ArgumentParser(description="Hic pipeline....")

    subparsers = parser.add_subparsers(help='sub-command help')
    
    parser_asm = subparsers.add_parser('asm', help='assemble genome')
    parser_asm.add_argument('-outdir', type=str, default="./", help='output directory', dest="outdir")
    parser_asm.add_argument('-threads', type=int, default=1, help='threads', dest ="threads")
    parser_asm.add_argument('-fq', type=str, required=True, help='long reads fastq', dest="fq")
    parser_asm.add_argument('-HiCfq', type=str, default=None, help='HiC fastq', dest="Hicfq")
    parser_asm.add_argument('-datatype', default="HiFi", help='HiFi or Nanopore', choices=["HiFi", "Nanopore"],dest="datatype")
    parser_asm.add_argument('-mem', default="100G", help="memory use", dest="Mem")
    parser_asm.set_defaults(func=assamble)
    
    parser_scf = subparsers.add_parser('scf', help='scaffold genome')
    parser_scf.add_argument('-ref', help="the reference genome for scaffold", dest = "ref")
    parser_scf.add_argument('-outdir', type=str, default="./", help='output directory', dest="outdir")
    parser_scf.add_argument('-fq1', help="fastq 1", dest="fq1")
    parser_scf.add_argument('-fq2', help="fastq 2", dest="fq2")
    parser_scf.add_argument('-threads', type=int, default=1, help='threads', dest ="threads")
    parser_scf.add_argument('-mt', type=str, default="bwa",choices=["bwa", "bowtie2", "chromap"], help='the mapping tools', dest ="mt")
    parser_scf.add_argument('-mem', default="100G", help="memory use", dest="Mem")
    parser_scf.add_argument('-agp', help="agpfiles", default=None, dest="agp")
    parser_scf.add_argument("-RE", help="the RE sites for scaffold",default="GATC", dest="RE")
    parser_scf.add_argument("-action", choices=["bam2juicer", "pair2juicer", "salsa2", "haphic"], help = "tools to use....", dest="action")
    parser_scf.add_argument("-bwaopt", help = "bwa mem options", default=" -SP5M ", dest="bwaopt")
    parser_scf.add_argument("-samtoolsopt", help="samtools view options", default=" -hF 268 ", dest="samtoolsopt")
    parser_scf.add_argument("-haphicopt", help="haphic options", default=" --correct_nrounds 2 -- max_inflation 5 ", dest="haphicopt")
    parser_scf.set_defaults(func=scaffold)
    
    parser_aln = subparsers.add_parser('align', help='alignment and stat')
    parser_aln.add_argument('-ref', help="the reference genome", dest = "ref")
    parser_aln.add_argument('-outdir', type=str, default="./", help='output directory', dest="outdir")
    parser_aln.add_argument('-fq1', help="fastq 1", dest="fq1")
    parser_aln.add_argument('-fq2', help="fastq 2", dest="fq2")
    parser_aln.add_argument('-threads', type=int, default=80, help='threads', dest ="threads")
    parser_aln.add_argument('-mem', default="200G", help="memory use", dest="Mem")
    parser_aln.add_argument("-datatype", choices=["ngs", "hifi", "ont"], help = "data type to use....", dest="datatype")
    parser_aln.set_defaults(func=alignment)
    
    args = parser.parse_args()
    args.func(args)
    

main()
    
    